from fastapi import APIRouter, Depends, FastAPI, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from fastapi.responses import PlainTextResponse, StreamingResponse
from typing import Annotated, Optional, Type
from os import environ
import logging
from enum import Enum
from .config import PaperAIConfig
from .core import PaperAICore
from .models import Project
from .interface.test import TestRequest, TestResponse
from .modules.block.sources import SourcesBlock, SourcesStartRequest, SourcesUpdateRequest
from .modules.block.topics import TopicsBlock, TopicsStartRequest, TopicsUpdateRequest
from .modules.block.review import ReviewBlock, ReviewStartRequest, ReviewUpdateRequest
from .modules.block.base import BaseBlock, ProgressChunk
from .modules.export.latex import LatexExportEngine


logging.basicConfig(level=logging.INFO)
config_path = environ.get('BACKEND_CONFIG_PATH', 'data/config.yaml')
config = PaperAIConfig.load(config_path)
core = PaperAICore(config)
core.start()


# Project token

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")


async def get_project_descriptor(token: Annotated[str, Depends(oauth2_scheme)]) -> Project:
    project = await core.project.get_descriptor(token)
    if project is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid project token",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return project


# Tags

class Tag(str, Enum):
    PROJECT = 'Project'
    ASSISTANT = 'AI Assistant'


# app

app = FastAPI(
    title='pAPer',
    openapi_tags=[
        {
            'name': Tag.PROJECT.value,
            'description': 'Project Descriptor, Token methods',
        },
        {
            'name': Tag.ASSISTANT.value,
            'description': 'AI Research Assistant for writing scientific papers',
        },
    ]
)

unauth_router = APIRouter()
auth_router = APIRouter(dependencies=[Depends(get_project_descriptor)])


app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_credentials=True,
    allow_methods=['*'],
    allow_headers=['*'],
)

# Project token

@unauth_router.post("/token", tags=[Tag.PROJECT.value])
async def project(form_data: Annotated[OAuth2PasswordRequestForm, Depends()]):
    """Make a new token.

    Does not check input data correctness.
    """
    project = await core.project.make()
    return {"access_token": project.token, "token_type": "bearer"}


@unauth_router.post("/project", tags=[Tag.PROJECT.value])
async def project():
    """Make a new project.

    Does not check input data correctness.
    """
    project = await core.project.make()
    return project


@auth_router.delete("/project/self", tags=[Tag.PROJECT.value])
async def project_self_delete(project: Annotated[Project, Depends(get_project_descriptor)]):
    """Delete the current project.
    """
    status = await core.project.delete(project)
    return {"status": status}


@unauth_router.delete("/project/token/{token}", tags=[Tag.PROJECT.value])
async def project_token_delete(token: str):
    """Delete project by token.
    """
    project = await core.project.get_descriptor(token)
    if project is None:
        raise HTTPException(status_code=404, detail="Project not found")
    status = await core.project.delete(project)
    return {"status": status}


# Assistant

@auth_router.post('/test', tags=[Tag.ASSISTANT.value])
async def test(project: Annotated[Project, Depends(get_project_descriptor)],
               request: TestRequest) -> TestResponse:
    """Test gigachat request.
    """
    return await core.test(project, request)


@auth_router.post("/idea", tags=[Tag.ASSISTANT.value])
async def idea(project: Annotated[Project, Depends(get_project_descriptor)],
               text: str,
               state_id: Optional[str] = None):
    """Post a new idea.
    """
    manager = core.state(project.token)
    current_state = await manager.get_state(state_id)
    data = await manager.get_data(current_state)
    data['idea'] = text
    new_state = await manager.append_state(current_state, data)
    return new_state


@auth_router.post("/reload", tags=[Tag.ASSISTANT.value])
async def reload(project: Annotated[Project, Depends(get_project_descriptor)],
                 state_id: Optional[str] = None):
    """Reload the state.
    """
    manager = core.state(project.token)
    current_state = await manager.get_state(state_id)
    return current_state


async def apply_block(project_token: str, state_id: str, block_cls: Type[BaseBlock], request):
    block_cls.config = config
    manager = core.state(project_token)
    current_state = await manager.get_state(state_id)
    data = await manager.get_data(current_state)
    block = block_cls.from_state(data, request)
    async def _wrap(gen):
        yield '['
        async for g in gen:
            if not g.done:
                yield g.model_dump_json() + ",\n"
                continue
            block.to_state(data)
            new_state = await manager.append_state(current_state, data)
            yield ProgressChunk(
                done=True,
                data=new_state.model_dump()
            ).model_dump_json()
        yield ']'
    return StreamingResponse(
        _wrap(block.run()),
        media_type='text/event-stream',
    )


@auth_router.post("/action/sources/start", tags=[Tag.ASSISTANT.value])
async def action_sources_start(project: Annotated[Project, Depends(get_project_descriptor)],
                               request: SourcesStartRequest):
    """Retrieve literature sources.
    """
    return await apply_block(project.token, request.state_id, SourcesBlock, request)


@auth_router.post("/action/sources/update", tags=[Tag.ASSISTANT.value])
async def action_sources_update(project: Annotated[Project, Depends(get_project_descriptor)],
                                request: SourcesUpdateRequest):
    """Update literature sources from client.
    """
    manager = core.state(project.token)
    current_state = await manager.get_state(request.state_id)
    data = await manager.get_data(current_state)
    data['sources'] = request.sources
    new_state = await manager.append_state(current_state, data)
    return new_state


@auth_router.post("/action/topics/start", tags=[Tag.ASSISTANT.value])
async def action_topics_start(project: Annotated[Project, Depends(get_project_descriptor)],
                              request: TopicsStartRequest):
    """Retrieve main topics.
    """
    return await apply_block(project.token, request.state_id, TopicsBlock, request)


@auth_router.post("/action/topics/update", tags=[Tag.ASSISTANT.value])
async def action_topics_update(project: Annotated[Project, Depends(get_project_descriptor)],
                               request: TopicsUpdateRequest):
    """Retrieve literature sources from client.
    """
    manager = core.state(project.token)
    current_state = await manager.get_state(request.state_id)
    data = await manager.get_data(current_state)
    data['topics'] = request.topics
    new_state = await manager.append_state(current_state, data)
    return new_state


@auth_router.post("/action/review/start", tags=[Tag.ASSISTANT.value])
async def action_review_start(project: Annotated[Project, Depends(get_project_descriptor)],
                              request: ReviewStartRequest):
    """Generate a literature review.
    """
    return await apply_block(project.token, request.state_id, ReviewBlock, request)


@auth_router.post("/action/review/update", tags=[Tag.ASSISTANT.value])
async def action_review_update(project: Annotated[Project, Depends(get_project_descriptor)],
                               request: ReviewUpdateRequest):
    """Manually update the literature review.
    """
    manager = core.state(project.token)
    current_state = await manager.get_state(request.state_id)
    data = await manager.get_data(current_state)
    data['review'] = request.review
    new_state = await manager.append_state(current_state, data)
    return new_state


@unauth_router.get("/export/latex/{project_token}/{state_token}", response_class=PlainTextResponse)
async def latex(project_token: str, state_token: str):
    if state_token == '' or state_token == 'recent':
        state_token = None
    project = await get_project_descriptor(project_token)
    if project is None:
        return f'Invalid project token: {project_token}'
    manager = core.state(project.token)
    current_state = await manager.get_state(state_token)
    data = await manager.get_data(current_state)
    return LatexExportEngine().export(data)


# include routers

app.include_router(unauth_router)
app.include_router(auth_router)
